{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "eeMrMI0-1Dhu"
   },
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "import numpy as np\n",
    "from IPython.display import clear_output\n",
    "from pydrake.all import (\n",
    "    AddMultibodyPlantSceneGraph,\n",
    "    DiagramBuilder,\n",
    "    JacobianWrtVariable,\n",
    "    MathematicalProgram,\n",
    "    MeshcatVisualizer,\n",
    "    PiecewisePolynomial,\n",
    "    Solve,\n",
    "    StartMeshcat,\n",
    ")\n",
    "\n",
    "from manipulation import running_as_notebook\n",
    "from manipulation.meshcat_utils import plot_mathematical_program\n",
    "from manipulation.scenarios import AddTwoLinkIiwa\n",
    "\n",
    "\n",
    "# This one is specific to this notebook, but I'm putting it in the header to make it less distracting.\n",
    "def Visualizer(MakeMathematicalProgram):\n",
    "    builder = DiagramBuilder()\n",
    "\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.0)\n",
    "    twolink = AddTwoLinkIiwa(plant, q0=[0.0, 0.0])\n",
    "    hand = plant.GetFrameByName(\"iiwa_link_ee\")\n",
    "    plant.Finalize()\n",
    "\n",
    "    MeshcatVisualizer.AddToBuilder(builder, scene_graph, meshcat)\n",
    "    diagram = builder.Build()\n",
    "    context = diagram.CreateDefaultContext()\n",
    "    plant_context = plant.GetMyContextFromRoot(context)\n",
    "\n",
    "    meshcat.Delete()\n",
    "    meshcat.SetProperty(\"/Background\", \"top_color\", [0, 0, 0])\n",
    "    meshcat.SetProperty(\"/Background\", \"bottom_color\", [0, 0, 0])\n",
    "    meshcat.SetProperty(\"/Grid\", \"visible\", False)\n",
    "\n",
    "    X, Y = np.meshgrid(np.linspace(-5, 5, 35), np.linspace(-5, 5, 31))\n",
    "\n",
    "    def visualize(q, v_Gdesired=[1.0, 0.0], t=None):\n",
    "        if t:\n",
    "            context.SetTime(t)\n",
    "        plant.SetPositions(plant_context, q)\n",
    "        diagram.ForcedPublish(context)\n",
    "\n",
    "        J_G = plant.CalcJacobianTranslationalVelocity(\n",
    "            plant_context,\n",
    "            JacobianWrtVariable.kQDot,\n",
    "            hand,\n",
    "            [0, 0, 0],\n",
    "            plant.world_frame(),\n",
    "            plant.world_frame(),\n",
    "        )\n",
    "        J_G = J_G[[0, 2], :]  # Ignore Y.\n",
    "        print(\"J_G = \")\n",
    "        print(\n",
    "            np.array2string(\n",
    "                J_G, formatter={\"float\": lambda x: \"{:5.2f}\".format(x)}\n",
    "            )\n",
    "        )\n",
    "\n",
    "        prog = MakeMathematicalProgram(q, J_G, v_Gdesired)\n",
    "        result = Solve(prog)\n",
    "        plot_mathematical_program(meshcat, \"QP\", prog, X, Y, result=result)\n",
    "        # TODO: Add set_object to meshcat.Animation\n",
    "        if False:  # meshcat._is_recording:\n",
    "            with meshcat._animation.at_frame(\n",
    "                v, meshcat._recording_frame_num\n",
    "            ) as m:\n",
    "                plot_mathematical_program(m, prog, X, Y, result=result)\n",
    "        clear_output(wait=True)\n",
    "\n",
    "    return visualize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the visualizer.\n",
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "hSeD3PBotJUU"
   },
   "source": [
    "# Differential Inverse Kinematics as a Quadratic Program\n",
    "\n",
    "## Define your mathematical program here.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "6F_KZoc1QDkY",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def MakeMathematicalProgram(q, J_G, v_Gdesired):\n",
    "    prog = MathematicalProgram()\n",
    "    v = prog.NewContinuousVariables(2, \"v\")\n",
    "    v_max = 3.0\n",
    "\n",
    "    error = J_G @ v - np.asarray(v_Gdesired)\n",
    "    prog.AddCost(error.dot(error))\n",
    "    prog.AddBoundingBoxConstraint(-v_max, v_max, v)\n",
    "\n",
    "    return prog"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "4i4xIYtmrIa-"
   },
   "source": [
    "## Visualize a particular joint angle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "CM9vi0TtrIa_",
    "tags": []
   },
   "outputs": [],
   "source": [
    "visualize = Visualizer(MakeMathematicalProgram)\n",
    "\n",
    "q = [-np.pi / 2.0 + 0.5, 1.0]\n",
    "v_Gdesired = [0.5, 0.0]\n",
    "visualize(q, v_Gdesired)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ttJ90nfLrIbD"
   },
   "source": [
    "## Animated joint trajectory (passing through the singularity)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ugsWyLNOrIbE",
    "tags": []
   },
   "outputs": [],
   "source": [
    "visualize = Visualizer(MakeMathematicalProgram)\n",
    "\n",
    "v_Gdesired = [1.0, 0.0]\n",
    "T = 2.0\n",
    "q = PiecewisePolynomial.FirstOrderHold(\n",
    "    [0, T, 2 * T],\n",
    "    np.array(\n",
    "        [\n",
    "            [-np.pi / 2.0 + 1.0, -np.pi / 2.0 - 1.0, -np.pi / 2.0 + 1.0],\n",
    "            [2.0, -2.0, 2],\n",
    "        ]\n",
    "    ),\n",
    ")\n",
    "\n",
    "nx = 35\n",
    "ny = 31\n",
    "X, Y = np.meshgrid(np.linspace(-5, 5, nx), np.linspace(-5, 5, ny))\n",
    "D = np.vstack((X.reshape(1, -1), Y.reshape(1, -1)))\n",
    "for i in range(2):\n",
    "    for t in np.linspace(0, 2 * T, num=100):\n",
    "        visualize(q.value(t), v_Gdesired, t=t)\n",
    "        if not running_as_notebook:\n",
    "            break\n",
    "        time.sleep(0.05)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "GPayCPx7rIbI"
   },
   "source": [
    "## Trajectory slider\n",
    "\n",
    "TODO(russt): I can remove this once I'm able to save the plotted surfaces in the meshcat animation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "KbPqBtd2rIbJ",
    "tags": []
   },
   "outputs": [],
   "source": [
    "visualize = Visualizer(MakeMathematicalProgram)\n",
    "\n",
    "v_Gdesired = [1.0, 0.0]\n",
    "T = 2.0\n",
    "qtraj = PiecewisePolynomial.FirstOrderHold(\n",
    "    [0, T], np.array([[-np.pi / 2.0 + 1.0, -np.pi / 2.0 - 1.0], [2.0, -2.0]])\n",
    ")\n",
    "visualize(qtraj.value(0), v_Gdesired)\n",
    "\n",
    "meshcat.AddSlider(\"time\", min=0, max=T, step=0.05, value=0)\n",
    "meshcat.AddButton(\"Stop Interaction Loop\")\n",
    "while meshcat.GetButtonClicks(\"Stop Interaction Loop\") < 1:\n",
    "    t = meshcat.GetSliderValue(\"time\")\n",
    "    visualize(qtraj.value(t), v_Gdesired)\n",
    "    if not running_as_notebook:\n",
    "        break\n",
    "    time.sleep(0.05)\n",
    "meshcat.DeleteAddedControls()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "DNhK1DCOrIbO"
   },
   "source": [
    "## Joint Sliders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Kn4igrtmrIbO",
    "tags": []
   },
   "outputs": [],
   "source": [
    "visualize = Visualizer(MakeMathematicalProgram)\n",
    "\n",
    "q = [-np.pi / 2.0 + 0.5, 1.0]\n",
    "v_Gdesired = [1.0, 0.0]\n",
    "visualize(q, v_Gdesired)\n",
    "\n",
    "\n",
    "def _q_callback(change, index):\n",
    "    q[index] = change.new\n",
    "    visualize(q, v_Gdesired)\n",
    "\n",
    "\n",
    "def _vG_callback(change, index):\n",
    "    v_Gdesired[index] = change.new\n",
    "    visualize(q, v_Gdesired)\n",
    "\n",
    "\n",
    "meshcat.AddSlider(\"q0\", value=q[0], min=-np.pi, max=np.pi, step=0.1)\n",
    "meshcat.AddSlider(\"q1\", value=q[1], min=-np.pi, max=np.pi, step=0.1)\n",
    "meshcat.AddSlider(\"v_G_W0\", value=v_Gdesired[0], min=-4, max=4, step=0.1)\n",
    "meshcat.AddSlider(\"v_G_W1\", value=v_Gdesired[1], min=-4, max=4, step=0.1)\n",
    "\n",
    "meshcat.AddButton(\"Stop Interaction Loop\")\n",
    "while meshcat.GetButtonClicks(\"Stop Interaction Loop\") < 1:\n",
    "    q = [meshcat.GetSliderValue(\"q0\"), meshcat.GetSliderValue(\"q1\")]\n",
    "    v_Gdesired = [\n",
    "        meshcat.GetSliderValue(\"v_G_W0\"),\n",
    "        meshcat.GetSliderValue(\"v_G_W1\"),\n",
    "    ]\n",
    "    visualize(q, v_Gdesired)\n",
    "    if not running_as_notebook:\n",
    "        break\n",
    "    time.sleep(0.05)\n",
    "meshcat.DeleteAddedControls()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "Robotic Manipulation - Lecture 5 Breakout.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3.10.6 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6"
  },
  "vscode": {
   "interpreter": {
    "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
